import os
import random
import warnings

import torch
import torch.backends.cudnn as cudnn
from torch import optim, nn, hub

from config import Config
from training import models
from training.datasets.oulu_npu_dataset import get_oulu_npu_dataloader
from training.models.compactnet import compactnet
from training.tools.train_utils import parse_args, train, validate

cudnn.benchmark = True

CONFIG = Config()
hub.set_dir(CONFIG['TORCH_HOME'])


def main():
    args = parse_args()
    print(args)

    if args.seed is not None:
        random.seed(args.seed)
        torch.manual_seed(args.seed)
        cudnn.deterministic = True
        warnings.warn('You have chosen to seed training.' +
                      'This will turn on the CUDNN deterministic setting, ' +
                      'which can slow down your training considerably! ' +
                      'You may see unexpected behavior when restarting from checkpoints.')

    # model = models.__dict__[args.arch](pretrained=True)
    model = compactnet(feature_extractor=models.__dict__[args.arch], pretrained=True)
    # model = discriminativecolorspace(feature_extractor=models.__dict__[args.arch], pretrained=True)

    train_loader, val_loader = get_oulu_npu_dataloader(model, args)

    if args.gpu is not None:
        torch.cuda.set_device(args.gpu)
        model = model.cuda(args.gpu)
    else:
        model = torch.nn.DataParallel(model).cuda()

    optimizer = optim.Adam(model.parameters(), lr=args.lr)
    # optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum)
    criterion = nn.CrossEntropyLoss().cuda()
    # criterion = TripletLoss().cuda()

    start_epoch = 1
    best_prec1 = 0.
    for epoch in range(start_epoch, args.epochs + 1):
        train(train_loader, model, optimizer, criterion, epoch, args)
        if epoch % 2 == 0 or epoch == args.epochs:
            acc1 = validate(val_loader, model, criterion, epoch, args)
            best_prec1 = max(acc1, best_prec1)
        save_model = ((epoch == args.epochs) or (epoch % 5 == 0))
        if save_model:
            print("Save model...")
            torch.save({
                'epoch': epoch,
                'arch': args.arch,
                'state_dict': model.state_dict(),
                'best_prec1': best_prec1,
                'optimizer': optimizer.state_dict(),
            }, os.path.join('weights', '{}_{}_{}.pt'.format(args.arch, args.prefix, epoch)))


if __name__ == '__main__':
    main()
